package io.github.heollhai.common.utils.tree;

import io.github.heollhai.common.utils.VerifyUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.ObjectUtils;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;

/**
 * @Author: lmf
 * @Create: 2024/8/7 10:47
 * Description:
 */


@Slf4j
public class TreeUtils {
    public static <T extends TreeNode<T, K, V>, K, V> List<T> getRoot(List<T> treeNode) {
        HashMap<K, T> nodeMap = new HashMap<>();
        for (T t : treeNode) {
            nodeMap.put(t.getId(), t);
        }

        // 找出所有的父节点 pid = 0 父节点 ->  0 || “0”
        List<T> rootNodes = new ArrayList<>();
        for (T node : treeNode) {
            if (node.getPid() instanceof Integer || node.getPid() instanceof String) {
                if (node.getPid().equals(0) || "0".equals(node.getPid())) {
                    rootNodes.add(node);
                }
            }
        }

        // 递归构建树结构
        for (T rootNode : rootNodes) {
            buildTree(rootNode, nodeMap);
        }
        // 对构建出来的树结构按照sort字段进行排序
        rootNodes =  sortTree(rootNodes);
        return rootNodes;
    }

    /**
     * 将树结构平铺为一维数组 （主要为了让部分先按照树的结构进行排序，这样在平铺出来的就是排序后的结构）
     */
    public static <T extends TreeNode<T, K, V>, K, V> List<T> flattenTree(List<T> treeNode) {
        List<T> flattenedTree = new ArrayList<>();
        for (T node : treeNode) {
            flattenedTree.add(node);
            if (VerifyUtils.isNotEmpty(node.getChildren())) {
                flattenedTree.addAll(flattenTree(node.getChildren()));
            }
        }
        return flattenedTree;
    }

    /**
     * 递归构建树结构
     */
    private static <T extends TreeNode<T, K, V>, K, V> void buildTree(T rootNode, HashMap<K, T> nodeMap) {
        for (T node : nodeMap.values()) {
            if (node.getPid().equals(rootNode.getId())) {
                rootNode.addChildren(node);
                buildTree(node, nodeMap);
            }
        }
    }

    /**
     * 对构建出来的树结构按照sort字段进行排序
     */
    private static <T extends TreeNode<T, K, V>, K, V> List<T>  sortTree(List<T> treeNode) {
        Comparator<T> comparator = Comparator.comparing(TreeNode::getSort);
        treeNode = treeNode.stream()
                .sorted(comparator.reversed())
                .toList();
        for (T node : treeNode) {
            if (VerifyUtils.isNotEmpty(node.getChildren())) {
               node.setChildren(sortTree(node.getChildren()));
            }
        }
        return treeNode;
    }

    /**
     * 根据id查询对应的所有子数据
     */
    public static <T extends TreeNode<T, K, V>, K, V> T findNodeById(List<T> treeRoot, V targetId) {

        for (T tree : treeRoot) {
            if (tree.getId().equals(targetId)) {
                return tree;
            }
            if (VerifyUtils.isNotEmpty(tree.getChildren())) {

                T result = findNodeById(tree.getChildren(), targetId);
                if(ObjectUtils.isNotEmpty(result)){
                    return result;
                }
            }
        }
        return null;
    }

}
